"""
optimal_knn_webapp_pinecone.py - 使用 Pinecone 的手写数字识别 Web 应用

修改自 IPR 问题 2: 使用最佳 KNN 模型识别手写数字的 APP 的 optimal_knn_webapp.py 文件
使其能够实现第一次推理速度较慢，但是第二次推理速度较快的效果

要求：
- 修改原有的 optimal_knn_webapp.py 文件
- 使用 Pinecone 云端向量数据库进行推理
- 第一次推理会较慢（需要连接 Pinecone）
- 第二次及之后的推理会更快（已建立连接）
"""

from __future__ import annotations

import io
import os
import time
from pathlib import Path
from typing import Tuple, Any, Optional

import numpy as np
from PIL import Image, ImageOps, ImageFilter

import gradio as gr
try:
    import cv2  # 可选：用于更快的 Otsu 与膨胀
except Exception:
    cv2 = None

try:
    from pinecone import Pinecone
    PINECONE_AVAILABLE = True
except ImportError:
    PINECONE_AVAILABLE = False
    print("警告: pinecone 未安装。请运行: pip install pinecone")

# Pinecone 配置（优先从环境变量读取，便于本地和部署时管理机密）
_DEFAULT_PINECONE_API_KEY = "pcsk_3q57Z8_E4VmVtFgcXQdmPDy2WkvM4vE33u2yQkjVegKobLBKy1s8KSAwQohW2SJiPFgpw7"
PINECONE_API_KEY = os.getenv("PINECONE_API_KEY", _DEFAULT_PINECONE_API_KEY)
INDEX_NAME = os.getenv("PINECONE_INDEX_NAME", "mnist-index")

# 全局变量：Pinecone 连接（第一次推理时初始化，之后复用）
_PINECONE_CLIENT = None
_PINECONE_INDEX = None
_FIRST_INFERENCE = True


def init_pinecone():
    """
    初始化 Pinecone 连接（仅在第一次推理时调用）
    """
    global _PINECONE_CLIENT, _PINECONE_INDEX, _FIRST_INFERENCE
    
    if not PINECONE_AVAILABLE:
        raise ImportError("请先安装 pinecone: pip install pinecone")
    if not PINECONE_API_KEY:
        raise RuntimeError("未检测到 PINECONE_API_KEY，请在系统环境变量中设置后重试。")
    
    if _PINECONE_CLIENT is None:
        print("正在初始化 Pinecone 连接...")
        _PINECONE_CLIENT = Pinecone(api_key=PINECONE_API_KEY)
        # 先检查索引是否存在，避免 404 错误
        try:
            existing_indexes = [idx.name for idx in _PINECONE_CLIENT.list_indexes()]
        except Exception as e:
            raise RuntimeError(f"获取 Pinecone 索引列表失败：{e}")

        if INDEX_NAME not in existing_indexes:
            raise RuntimeError(
                f"Pinecone 索引 '{INDEX_NAME}' 不存在。\n"
                f"请先运行 'pinecone_train.py' 创建并上传数据，或设置环境变量 PINECONE_INDEX_NAME 指向已存在的索引。"
            )

        _PINECONE_INDEX = _PINECONE_CLIENT.Index(INDEX_NAME)
        print("Pinecone 连接成功！")
        _FIRST_INFERENCE = False


def _to_pil(input_obj: Any) -> Image.Image:
    """将 Sketchpad 可能返回的多种类型(dict/ndarray/PIL)统一转为 PIL.Image.

    兼容格式：
    - PIL.Image
    - numpy.ndarray (H,W) 或 (H,W,3|4)
    - dict：常见键 image/composite/background/layers（取首个可解析的图层）
    """
    if input_obj is None:
        raise ValueError("未检测到图像")

    # 已是 PIL
    if isinstance(input_obj, Image.Image):
        pil = input_obj
    # numpy 数组
    elif isinstance(input_obj, np.ndarray):
        pil = Image.fromarray(input_obj)
    # gradio 某些版本的 Sketchpad/ImageEditor 可能返回 dict
    elif isinstance(input_obj, dict):
        candidate = None
        for key in ("image", "composite", "background"):
            if key in input_obj and input_obj[key] is not None:
                candidate = input_obj[key]
                break
        # 尝试 layers（若存在），取最后一层
        if candidate is None and isinstance(input_obj.get("layers"), list) and input_obj["layers"]:
            last = input_obj["layers"][-1]
            if isinstance(last, dict):
                candidate = last.get("image") or last.get("composite")
            else:
                candidate = last

        if candidate is None:
            raise ValueError("无法从输入字典解析图像内容")

        if isinstance(candidate, Image.Image):
            pil = candidate
        elif isinstance(candidate, np.ndarray):
            pil = Image.fromarray(candidate)
        else:
            # 有些实现可能是 base64/bytes，这里先简单兜底
            raise ValueError("不支持的图像数据类型: " + type(candidate).__name__)
    else:
        raise ValueError("不支持的输入类型: " + type(input_obj).__name__)

    # 去除 alpha：合成到白底，避免反相后出现暗底
    if pil.mode in ("RGBA", "LA"):
        bg = Image.new("RGBA", pil.size, (255, 255, 255, 255))
        pil = Image.alpha_composite(bg, pil.convert("RGBA")).convert("RGB")

    return pil


def _otsu_threshold(arr: np.ndarray) -> int:
    """简易 Otsu 自动阈值计算，返回 0..255 的整型阈值。"""
    # 优先使用 OpenCV 的 Otsu（C 实现更快）
    if cv2 is not None:
        try:
            _arr = arr.astype(np.uint8, copy=False)
            _, th = cv2.threshold(_arr, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
            return int(_)
        except Exception:
            pass
    hist, _ = np.histogram(arr, bins=256, range=(0, 255))
    total = arr.size
    sum_total = np.dot(np.arange(256), hist)
    sumB = 0.0
    wB = 0.0
    max_var = -1.0
    thresh = 0
    for t in range(256):
        wB += hist[t]
        if wB == 0:
            continue
        wF = total - wB
        if wF == 0:
            break
        sumB += t * hist[t]
        mB = sumB / wB
        mF = (sum_total - sumB) / wF
        varBetween = wB * wF * (mB - mF) ** 2
        if varBetween > max_var:
            max_var = varBetween
            thresh = t
    return int(thresh)


def _preprocess_to_digits(
    img: Image.Image,
    invert: bool = True,
    thresh: int = 10,
    padding: int = 4,
    blur_radius: float = 1.0,
    dilate: int = 0,
    auto_thresh: bool = False,
) -> np.ndarray:
    """
    将任意输入图像预处理为 sklearn digits (8x8, 值域 0..16) 的扁平 1x64 向量。

    步骤：转灰度 -> 反相(使数字为亮、背景暗) -> 缩放到 8x8 -> 归一化到 0..16 -> flatten
    """
    # 1. 转灰度
    gray = ImageOps.grayscale(img)
    arr = np.array(gray, dtype=np.float32)

    # 2. 反相（可选）
    if invert:
        arr = 255.0 - arr

    # 3. 自动阈值（可选）
    if auto_thresh:
        thresh = _otsu_threshold(arr.astype(np.uint8))

    # 4. 二值化 + 找前景
    binary = (arr > thresh).astype(np.uint8) * 255

    # 5. 膨胀（可选）
    if dilate > 0:
        if cv2 is not None:
            kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
            for _ in range(dilate):
                binary = cv2.dilate(binary, kernel)
        else:
            # 简易膨胀：最大池化
            from scipy.ndimage import maximum_filter
            for _ in range(dilate):
                binary = maximum_filter(binary, size=3)

    # 6. 找外接框
    coords = np.column_stack(np.where(binary > 0))
    if coords.size == 0:
        # 空白图像，返回全零向量
        return np.zeros(64, dtype=np.float32)

    y_min, x_min = coords.min(axis=0)
    y_max, x_max = coords.max(axis=0)

    # 7. 裁剪 + padding
    y_min = max(0, y_min - padding)
    x_min = max(0, x_min - padding)
    y_max = min(binary.shape[0], y_max + padding + 1)
    x_max = min(binary.shape[1], x_max + padding + 1)

    cropped = binary[y_min:y_max, x_min:x_max]

    # 8. 缩放到 8x8
    pil_crop = Image.fromarray(cropped, mode="L")
    if blur_radius > 0:
        pil_crop = pil_crop.filter(ImageFilter.GaussianBlur(radius=blur_radius))
    pil_8x8 = pil_crop.resize((8, 8), Image.LANCZOS)

    # 9. 归一化到 0..16
    arr_8x8 = np.array(pil_8x8, dtype=np.float32)
    arr_8x8 = arr_8x8 / 255.0 * 16.0

    # 10. flatten
    return arr_8x8.flatten()


def predict_digit_with_pinecone(
    input_img: Any,
    invert: bool = True,
    thresh: int = 10,
    padding: int = 4,
    blur_radius: float = 1.0,
    dilate: int = 0,
    preview_size: int = 128,
    show_grid: bool = False,
    auto_thresh: bool = True,
    fast_mode: bool = False,
    k: int = 11,
) -> Tuple[int, str, Optional[Image.Image]]:
    """
    使用 Pinecone 进行预测
    
    第一次调用会较慢（需要初始化 Pinecone 连接）
    之后的调用会更快（复用已建立的连接）
    """
    global _FIRST_INFERENCE
    
    try:
        # 第一次推理：初始化 Pinecone
        start_time = time.time()
        
        if _PINECONE_INDEX is None:
            init_pinecone()
            init_info = "（首次推理，正在连接 Pinecone...）"
        else:
            init_info = "（使用已建立的 Pinecone 连接）"
        
        if input_img is None:
            return -1, "请在画板中写一个数字", None

        # 预处理图像
        pil_img = _to_pil(input_img)
        if fast_mode:
            features = _preprocess_to_digits(
                pil_img,
                invert=invert,
                thresh=thresh,
                padding=padding,
                blur_radius=0.0,
                dilate=0,
                auto_thresh=False,
            )
        else:
            features = _preprocess_to_digits(
                pil_img,
                invert=invert,
                thresh=thresh,
                padding=padding,
                blur_radius=blur_radius,
                dilate=dilate,
                auto_thresh=auto_thresh,
            )
        
        # 使用 Pinecone 查询
        query_start = time.time()
        results = _PINECONE_INDEX.query(
            vector=features.tolist(),
            top_k=k,
            include_metadata=True
        )
        query_time = time.time() - query_start
        
        matches = (results or {}).get('matches') or []
        # 生成 8x8 预处理可视化图（无论是否查询到结果，都用于展示）
        img8 = (features.reshape(8, 8) / 16.0 * 255.0).astype(np.uint8)
        pv = 128 if fast_mode else max(32, int(preview_size))
        viz = Image.fromarray(img8, mode="L").resize((pv,) * 2, Image.NEAREST)

        if not matches:
            total_time = time.time() - start_time
            info = f"未能识别，请重新书写\n用时: {total_time:.2f}秒"
            if show_grid:
                viz = viz.convert("RGB")
                import PIL.ImageDraw as ImageDraw
                draw = ImageDraw.Draw(viz)
                step = viz.size[0] // 8
                for i in range(1, 8):
                    draw.line([(i * step, 0), (i * step, viz.size[1])], fill=(200, 200, 200), width=1)
                    draw.line([(0, i * step), (viz.size[0], i * step)], fill=(200, 200, 200), width=1)
            return -1, info, viz

        # 获取最近邻的标签并投票
        labels = [int(match['metadata']['label']) for match in matches]
        # 修复 NumPy 2.1 deprecation warning：确保 labels 是整数数组
        labels = np.array(labels, dtype=int)
        predicted_label = int(np.bincount(labels).argmax())
        
        # 计算总时间
        total_time = time.time() - start_time
        
        # 生成说明信息
        info = f"识别结果: {predicted_label}\n"
        info += f"用时: {total_time:.2f}秒"
        
        if show_grid:
            # 叠加网格，帮助观察每个像素格
            viz = viz.convert("RGB")
            import PIL.ImageDraw as ImageDraw
            draw = ImageDraw.Draw(viz)
            step = viz.size[0] // 8
            for i in range(1, 8):
                draw.line([(i * step, 0), (i * step, viz.size[1])], fill=(200, 200, 200), width=1)
                draw.line([(0, i * step), (viz.size[0], i * step)], fill=(200, 200, 200), width=1)
        
        return predicted_label, info, viz
        
    except Exception as e:
        return -1, f"识别出错: {e}", None


# 创建 Gradio 接口
with gr.Blocks(title="手写数字识别") as demo:
    gr.Markdown("""
    # 手写数字识别
    
    在画板中写下一个数字，系统将给出识别结果。
    """)

    with gr.Row():
        sketch = gr.Sketchpad(
            label="画板",
            type="numpy",
            height=280,
            width=280,
        )
        with gr.Column():
            gr.Markdown("### 📝 使用提示")
            gr.Markdown("• 写粗一点，数字要饱满\n• 数字要居中，占据大部分画板")
            
            with gr.Accordion("高级设置", open=False):
                invert_chk = gr.Checkbox(value=True, label="反相处理")
                with gr.Row():
                    thresh_sld = gr.Slider(0, 255, value=50, step=1, label="阈值")
                    pad_sld = gr.Slider(0, 32, value=8, step=1, label="边距")
                with gr.Row():
                    blur_sld = gr.Slider(0, 3, value=1.0, step=0.5, label="模糊")
                    dilate_sld = gr.Slider(0, 5, value=2, step=1, label="加粗")
                with gr.Row():
                    prev_sld = gr.Slider(64, 256, value=192, step=32, label="预览尺寸")
                    grid_chk = gr.Checkbox(value=True, label="显示网格")
                with gr.Row():
                    k_sld = gr.Slider(1, 20, value=11, step=1, label="K 值")
                auto_thresh_chk = gr.Checkbox(value=True, label="自动阈值")
                fast_chk = gr.Checkbox(value=False, label="快速模式")
            
            btn = gr.Button("识别", variant="primary", size="lg")
            out_label = gr.Number(label="识别结果", precision=0)
            out_text = gr.Textbox(label="详细信息", lines=2)
            out_preview = gr.Image(label="处理预览", type="pil")

            # 清除按钮
            try:
                gr.ClearButton(
                    components=[sketch, out_preview],
                    value="清除画板",
                )
            except Exception:
                clear_btn = gr.Button("清除画板")
                def _clear_only_canvas_preview():
                    return None, None
                clear_btn.click(
                    _clear_only_canvas_preview,
                    inputs=None,
                    outputs=[sketch, out_preview],
                )

    btn.click(
        predict_digit_with_pinecone,
        inputs=[sketch, invert_chk, thresh_sld, pad_sld, blur_sld, dilate_sld, 
                prev_sld, grid_chk, auto_thresh_chk, fast_chk, k_sld],
        outputs=[out_label, out_text, out_preview],
    )


if __name__ == "__main__":
    print("手写数字识别系统启动中...")
    demo.launch()
